# -*- coding: utf-8 -*-

import configparser
import logging
import os
import sys

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd


config = configparser.ConfigParser()
config_file = 'config.ini'
if not os.path.exists(config_file):
    raise FileNotFoundError(f"Configuration file '{config_file}' not found.")


config.read(config_file)
DATA_DIR = config['DEFAULT'].get('DataDir', 'd:/Python/study_data')
LIST_DIR = config['DEFAULT'].get('List', 'list')
LOG_FILE = config['DEFAULT'].get('LogFile_Podifan', 'log_podifan.txt')

if not os.path.exists(DATA_DIR):
    os.makedirs(DATA_DIR)

filelog = True
logger = logging.getLogger('log')
logger.setLevel(logging.DEBUG)
while logger.hasHandlers():
    for i in logger.handlers:
        logger.removeHandler(i)
# file log
# formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
formatter = logging.Formatter('%(message)s')

if filelog:
    fh = logging.FileHandler(LOG_FILE, encoding='utf-8')
    fh.setLevel(logging.DEBUG)
    fh.setFormatter(formatter)
    logger.addHandler(fh)

# console log
# formatter = logging.Formatter('%(message)s')
# ch = logging.StreamHandler(sys.stdout)
# ch.setLevel(logging.DEBUG)
# ch.setFormatter(formatter)
# logger.addHandler(ch)


def test_all_stocks(folder):
    hold_5_days = []
    hold_10_days = []
    hold_20_days = []
    hold_30_days = []
    hold_60_days = []
    for file in os.listdir(folder):
        if file.endswith('.csv'):
            path = os.path.join(folder, file)
            result = test_stock(path)
            hold_5_days.append(result[0])
            hold_10_days.append(result[1])
            hold_20_days.append(result[2])
            hold_30_days.append(result[3])
            hold_60_days.append(result[4])

    print(f"Average return for holding 5 days: {np.mean(hold_5_days):.4f}")
    print(f"Average return for holding 10 days: {np.mean(hold_10_days):.4f}")
    print(f"Average return for holding 20 days: {np.mean(hold_20_days):.4f}")
    print(f"Average return for holding 30 days: {np.mean(hold_30_days):.4f}")
    print(f"Average return for holding 60 days: {np.mean(hold_60_days):.4f}")


def test_stock(path, name):
    """
    Parameters:
    path (str): The file path to the stock data CSV file.

    Returns:
    None
    """
    # 获取股票数据
    data = pd.read_csv(path, parse_dates=['date'])

    # 获取股票数据
    data = pd.read_csv(path)

    # 确保索引是日期时间类型
    if 'date' in data.columns:
        data.index = pd.to_datetime(data.date)
    else:
        raise ValueError("The 'date' column is missing from the data.")

    # 确保索引是日期时间类型
    data.index = pd.to_datetime(data.date)

    # 计算5日均线
    data['5_day_mavg'] = data['close'].rolling(window=5).mean()

    # 计算10日均线
    data['10_day_mavg'] = data['close'].rolling(window=10).mean()

    # 判断下跌趋势
    data['20_day_mavg'] = data['close'].rolling(window=20).mean()
    data['falling_trend'] = (
            (data['close'] < data['20_day_mavg']) &
            (data['close'].shift(1) < data['20_day_mavg'].shift(1)) &
            (data['close'].shift(2) < data['20_day_mavg'].shift(2)) &
            (data['close'].shift(3) < data['20_day_mavg'].shift(3)) &
            (data['close'].shift(4) < data['20_day_mavg'].shift(4)) &
            (data['close'].shift(5) < data['20_day_mavg'].shift(5)) &
            (data['close'].shift(6) < data['20_day_mavg'].shift(6)) &
            (data['close'].shift(7) < data['20_day_mavg'].shift(7)) &
            (data['close'].shift(8) < data['20_day_mavg'].shift(8)) &
            (data['close'].shift(9) < data['20_day_mavg'].shift(9)) &
            (data['close'].shift(10) < data['20_day_mavg'].shift(10)) &
            (data['5_day_mavg'] < data['10_day_mavg']) &
            (data['5_day_mavg'].shift(1) < data['10_day_mavg'].shift(1)) &
            (data['5_day_mavg'].shift(2) < data['10_day_mavg'].shift(2)) &
            (data['5_day_mavg'].shift(3) < data['10_day_mavg'].shift(3)) &
            (data['5_day_mavg'].shift(4) < data['10_day_mavg'].shift(4)) &
            (data['5_day_mavg'].shift(5) < data['10_day_mavg'].shift(5)) &
            (data['5_day_mavg'].shift(6) < data['10_day_mavg'].shift(6)) &
            (data['5_day_mavg'].shift(7) < data['10_day_mavg'].shift(7)) &
            (data['5_day_mavg'].shift(8) < data['10_day_mavg'].shift(8)) &
            (data['5_day_mavg'].shift(9) < data['10_day_mavg'].shift(9)) &
            (data['5_day_mavg'].shift(10) < data['10_day_mavg'].shift(10)) &
            (data['10_day_mavg'] < data['20_day_mavg']) &
            (data['10_day_mavg'].shift(1) < data['20_day_mavg'].shift(1)) &
            (data['10_day_mavg'].shift(2) < data['20_day_mavg'].shift(2)) &
            (data['10_day_mavg'].shift(3) < data['20_day_mavg'].shift(3)) &
            (data['10_day_mavg'].shift(4) < data['20_day_mavg'].shift(4)) &
            (data['10_day_mavg'].shift(5) < data['20_day_mavg'].shift(5)) &
            (data['10_day_mavg'].shift(6) < data['20_day_mavg'].shift(6)) &
            (data['10_day_mavg'].shift(7) < data['20_day_mavg'].shift(7)) &
            (data['10_day_mavg'].shift(8) < data['20_day_mavg'].shift(8)) &
            (data['10_day_mavg'].shift(9) < data['20_day_mavg'].shift(9)) &
            (data['10_day_mavg'].shift(10) < data['20_day_mavg'].shift(10))

    )

    # 寻找放量阳线
    data['20_day_avg_volume'] = data['volume'].rolling(window=20).mean()
    data['bullish_volume'] = (data['volume'] > 1 * data['20_day_avg_volume']) & (data['close'] > data['5_day_mavg']) & (data['close'].pct_change() > 0.065)

    # 生成买入信号
    data['close_to_10_day_mavg'] = np.isclose(data['close'], data['10_day_mavg'], atol=0.04 * data['10_day_mavg'])

    data['buy_signal'] = (
            (data['close'].shift(5) < data['20_day_mavg'].shift(5)) & (data['falling_trend'].shift(5) & data['bullish_volume'].shift(4, fill_value=False) & (data['close'].shift(3) > data['10_day_mavg'].shift(3)) & data['close_to_10_day_mavg']) |
            (data['close'].shift(6) < data['20_day_mavg'].shift(6)) & (data['falling_trend'].shift(6) & data['bullish_volume'].shift(5, fill_value=False) & (data['close'].shift(4) > data['10_day_mavg'].shift(4)) & data['close_to_10_day_mavg']) |
            (data['close'].shift(7) < data['20_day_mavg'].shift(7)) & (data['falling_trend'].shift(7) & data['bullish_volume'].shift(6, fill_value=False) & (data['close'].shift(5) > data['10_day_mavg'].shift(5)) & data['close_to_10_day_mavg']) |
            (data['close'].shift(8) < data['20_day_mavg'].shift(8)) & (data['falling_trend'].shift(8) & data['bullish_volume'].shift(7, fill_value=False) & (data['close'].shift(6) > data['10_day_mavg'].shift(6)) & data['close_to_10_day_mavg']) |
            (data['close'].shift(9) < data['20_day_mavg'].shift(9)) & (data['falling_trend'].shift(9) & data['bullish_volume'].shift(8, fill_value=False) & (data['close'].shift(7) > data['10_day_mavg'].shift(7)) & data['close_to_10_day_mavg'])
    )
    data['buy_signal_'] = (data['buy_signal'] & ~(data['buy_signal'].shift(1) | data['buy_signal'].shift(2) | data['buy_signal'].shift(3) | data['buy_signal'].shift(4) | data['buy_signal'].shift(5)))
    # 检查生成的买入信号日期
    buy_signal_dates = data[data['buy_signal_']].index
    # print(f"Buy signal dates: {buy_signal_dates}")
    dict = {}
    dict[os.path.basename(path).split('.')[0]] = [str(date)[:10] for date in buy_signal_dates]
    # print(dict)
    return dict

    # 初始化用于记录不同持有期收益的DataFrame
    hold_periods = [5, 10, 20, 30, 60]
    hold_days = [f'hold_{period}_days' for period in hold_periods]
    returns_df = pd.DataFrame(index=data.index, columns=hold_days)

    # 计算不同持有期的收益
    for period in hold_periods:
        for buy_date in buy_signal_dates:
            sell_date_idx = data.index.get_loc(buy_date) + period
            if sell_date_idx < len(data.index):
                buy_price = data.loc[buy_date, 'close']
                sell_date = data.index[sell_date_idx]
                sell_price = data.loc[sell_date, 'close']
                if pd.isna(sell_price):
                    logging.warning(f"Sell price is NaN for buy date {buy_date} and sell date {sell_date}")
                    returns_df.loc[buy_date, f'hold_{period}_days'] = np.nan
                else:
                    returns_df.loc[buy_date, f'hold_{period}_days'] = (sell_price - buy_price) / buy_price
            else:
                # logging.warning(f"Sell date index {sell_date_idx} is out of range for buy date {buy_date}")
                returns_df.loc[buy_date, f'hold_{period}_days'] = np.nan

    # 绘制不同持有期的累计收益曲线
    results = []
    for period in hold_periods:
        cumulative_returns = (1 + returns_df[f'hold_{period}_days']).cumprod(skipna=True) - 1
        plt.plot(cumulative_returns, label=f'Hold {period} days')
        avg_return = returns_df[f'hold_{period}_days'].mean(skipna=True)
        results.append(avg_return)


def find_all_buys(path):
    all_buys = {}
    list_df = pd.read_csv(f'{LIST_DIR}/stocks.csv')
    stocks = dict(zip(list_df['名称'], list_df['代码']))
    for name, code in stocks.items():
        if 'ST' not in name:
            try:
                buys = test_stock(f'{DATA_DIR}/{path}/{code}.csv', name)
                for key, value in buys.items():
                    for date in value:
                        if date in all_buys:
                            all_buys[date].append(key)
                            # all_buys[date].append(name)
                        else:
                            all_buys[date] = [key]
                            # all_buys[date] = [name]
            except FileNotFoundError:
                # print(f"File not found for stock {name} with code {code}")
                continue

    all_buys_df = pd.DataFrame.from_dict(all_buys, orient='index').sort_index()
    all_buys_df.to_csv('all_buys_10_day_line.csv')

    # all_buys = sorted(all_buys.items())
    # with open('all_buys.txt', 'w', encoding='utf-8') as f:
    #     for date, stocks in all_buys:
    #         f.write(f'{date},{",".join(stocks)}\n')


if __name__ == '__main__':
    find_all_buys('stocks')
